# -*- coding: utf-8 -*-
"""
Created on Mon Dec 25 13:56:27 2023

@author: HD
"""

import numpy as np
import pandas as pd
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.statespace.sarimax import SARIMAX


import random
import ast


def arima_forecast(data, plan, f_plan, use_plan, order=(1, 1, 1), forecast_steps=3):
    """
    使用ARIMA模型对多个时间序列进行预测。

    参数：
    - data: 输入的时间序列数据，维度为（n_timesteps, n_samples）。
    - plan: 输入的发货计划序列数据，维度为（n_timesteps, n_samples），要求与data维度一样。
    - use_plan: 是否使用发货计划数据预测
    - order: ARIMA模型的阶数，形如 (p, d, q)。
    - forecast_steps: 预测的时间步数量。

    返回：
    - 预测结果，维度为（forecast_steps，n_samples）。
    """

    # 存储预测结果
    predictions = []

    # 对每个时间序列进行训练和预测
    for i in range(data.shape[1]):
        # 训练ARIMA模型
        if use_plan == 'False':
            model = ARIMA(data[:, i], order=order)
            model_fit = model.fit()
            forecast = model_fit.get_forecast(steps=forecast_steps)
            forecast_mean = forecast.predicted_mean
            predictions.append(forecast_mean)

        else:
            model = ARIMA(data[:, i], exog=plan[:, i], order=order)
            model_fit = model.fit()
            forecast = model_fit.get_forecast(steps=forecast_steps, exog=f_plan[:, i])
            forecast_mean = forecast.predicted_mean
            predictions.append(forecast_mean)

        # 进行未来预测
        #forecast = model_fit.get_forecast(steps=forecast_steps)

        # 获取预测结果
        #forecast_mean = forecast.predicted_mean
        #predictions.append(forecast_mean)

    # 将预测结果转换为数组，维度为（forecast_steps，n_samples）
    predictions = np.array(predictions).T
    #predictions = pd.DataFrame(predictions)
    predictions[predictions < 0] = 0
    return predictions


def forecast_with_autoarima(data, plan, f_plan, use_plan, order,s, d_range, D_range, p_range, q_range, P_range, Q_range, forecast_steps):
    """
    使用AutoARIMA模型进行预测

    参数：
    - df: 输入的DataFrame，包含时间序列数据，形状为(N, F)，N代表时间步，F代表不同的时间序列
    - plan: 输入的发货计划序列数据，维度为（n_timesteps, n_samples），要求与data维度一样。
    - use_plan: 是否使用发货计划数据预测
    - s: 季节性周期长度
    - d_range: 差分阶数的范围，例如(0, 1)表示可选0或1阶差分
    - D_range: 季节性差分阶数的范围，例如(0, 1)表示可选0或1阶季节性差分
    - p_range: 自回归阶数的范围，例如(0, 2, 4)表示可选0、1或2阶自回归
    - q_range: 移动平均阶数的范围，例如(0, 2, 4)表示可选0、1或2阶移动平均
    - P_range: 季节性自回归阶数的范围，例如range(2)表示可选0或1阶季节性自回归
    - Q_range: 季节性移动平均阶数的范围，例如range(2)表示可选0或1阶季节性移动平均
    - forecast_steps: 预测的时间步数

    返回值：
    - fc: 预测结果，形状为(F, forecast_steps)，每列包含对应时间序列的未来预测值
    """
    # import cudf
    from cuml.tsa.arima import ARIMA
    from cuml.tsa.auto_arima import AutoARIMA
    from cuml import ExponentialSmoothing
    import cudf
    import cupy as cp
    import numpy as np

    if use_plan == 'False':
        model = AutoARIMA(data)
        # 搜索最佳模型参数
        model.search(d=d_range, D=D_range, p=p_range, q=q_range,
                 P=P_range, Q=Q_range, method="css", truncate=20)
        model.fit(method="css-ml")
        fc = model.forecast(forecast_steps)
        fc_list = fc[0].tolist()

        # 取最后五行
        last_five_rows = data[-5:]
        # 计算每列的平均值
        average_values = np.mean(last_five_rows, axis=0)
        average_values_list = average_values.tolist()

        df = cudf.DataFrame.from_records(data)
        list = []
        np.set_printoptions(suppress=True)
        for col in df.columns:
            series_col = df[col]
            # series_col 现在是一个 cuDF Series，可以进行进一步的操作
            series_col_float = series_col.astype('float64')

            cu_hw = ExponentialSmoothing(series_col_float)
            cu_hw.fit()
            cu_pred = cu_hw.forecast(1)
            list.append(0 if cu_pred.values.item() < 0 else cu_pred.values.item())
        # list = [int(num) for num in list]
        min_list = [min(pair) for pair in zip(fc_list, list, average_values_list)]
        cupy_array = cp.array([min_list])
        # cupy_array = cp.array(list)
        # numpy_array = cupy_array.get()
        # print(cupy_array)
    else:   # 如果要用发货计划数据, 则不能使用AutoArima
        #model = AutoARIMA()
        #model.fit(data, exogenous=plan)
        #forecast = model.predict(n_periods=forecast_steps, exogenous= f_plan)
        #model = ARIMA(data, exog=)
        noise_level = 1e-6
        data += np.random.uniform(0, noise_level, size=(data.shape[0], data.shape[1]))
        model = ARIMA(data,exog=plan,order= (1,0,0)) #exog=plan
        model.fit()
        fc = model.forecast(forecast_steps,exog=f_plan)#exog=f_plan
    # 拟合模型
    #model.fit(method="css-ml")
    # 进行未来预测
    #fc = model.forecast(forecast_steps)
    return cupy_array

def forecast(data, plan, f_plan, forecast_steps=3, use_plan=False, use_gpu=False):
    """
    对多个时间序列进行预测，可选择使用ARIMA模型或AutoARIMA模型。

    参数：
    - data: 输入的时间序列数据，维度为（n_timesteps, n_samples）。
    - plan: 输入的发货计划序列数据，维度为（n_timesteps, n_samples），要求与data维度一样。
    - order: ARIMA模型的阶数，形如 (p, d, q)。
    - forecast_steps: 预测的时间步数量。
    - use_plan: 是否使用发货计划数据预测
    - use_gpu: 是否使用AutoARIMA模型。

    返回：
    - 预测结果，维度为（forecast_steps，n_samples）。
    """
    if use_gpu == 'True':
        order = (1, 0, 0)
        forecast = forecast_with_autoarima(data, plan, f_plan, use_plan, order,s=2, d_range=(0, 5), D_range=(0, 5),
                                           p_range=(0, 5), q_range=(0, 5),
                                           P_range=(0, 5), Q_range=(0, 5),
                                           forecast_steps=forecast_steps)
        print(forecast)  # 输出格式为 (forecast_steps, num_variables)  forecast_steps为未来预测的步数，num_variables为时间序列的个数
        return forecast
    else:
        order = (1, 0, 0)  # 无GPU情况下可以调整的参数,备选（1,0,0）（0,1,1）
        return arima_forecast(data, plan, f_plan = f_plan, use_plan=use_plan, order=order, forecast_steps=forecast_steps)


'''
# 示例用法
# 定义变量数量和观测值数量
file_path = 'F:/Project/RUL_GAN/yc231215.txt'
with open(file_path, 'r') as file:
    # 读取文件内容
    file_content = file.read()
# 将字符串转换为列表
try:
    my_list = ast.literal_eval(file_content)
    print("成功读取列表:", my_list)
except ValueError as e:
    print(f"发生错误: {e}")

my_list = np.array(my_list)

num_variables = 100  # 时间序列的个数
num_observations = 30   # 每个时间序列有num_observations个时间步
data = {f"Variable_{i+1}": [random.uniform(0,10000) for _ in range(num_observations)] for i in range(num_variables)}
df = pd.DataFrame(data)
'''
#
# # 读取数据
# file_path = 'yc231215.txt'
# with open(file_path, 'r') as file:
#     # 读取文件内容
#     file_content = file.read()
# # 将字符串转换为列表
# try:
#     my_list = ast.literal_eval(file_content)
#     print("成功读取列表:")
# except ValueError as e:
#     print(f"发生错误: {e}")
#
# #num_variables = 10  # 时间序列的个数
# # = 30   # 每个时间序列有num_observations个时间步
# #data = {f"Variable_{i+1}": [random.uniform(0,10000) for _ in range(num_observations)] for i in range(num_variables)}
#
# forecast_steps = 1  #未来预测的天数
# data = np.array(my_list)
# #data = np.random.rand(90, 7000)
# # h_plan = np.random.rand(90, 7000)
# h_plan = np.array(my_list)  #历史发货计划 格式与data相同
# f_plan = np.random.rand(forecast_steps, data.shape[1])#未来的发货计划，格式为 （forecast_steps，货物种类），如需要预测未来3天，就需要未来3天每个货物的发货计划
# is_gpu = 'True'  # 如有GPU，参数为True，无GPU则为False
# use_plan = 'False' # 参数为True，表示考虑发货计划，参数为False，表示不考虑发货计划
#
# forecast_result = forecast(data, plan=h_plan, f_plan =f_plan,
#                            forecast_steps=forecast_steps,
#                            use_gpu=is_gpu, use_plan = use_plan)
#
#
#
# print(forecast_result)



